///////////////////////////////////////////////////////////////////////////////////////////////////
/// \file formic/psi3/interface.cpp
///
/// \brief   contains definitions of PSI3 access functions
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#include <iostream>
#include <algorithm>

#include <formic/exception.h>
#include <formic/input/base.h>
#include <formic/lapack/interface.h>
#include <formic/psi3/interface.h>
#include <formic/ham/tei.h>

#ifdef FORMIC_HAVE_PSI3

// include headers needed for PSI3
extern "C" {
#include <stdio.h>
#include <stdlib.h>
#include <libciomr/libciomr.h>
#include <psifiles.h>
#include <libchkpt/chkpt.h>
#include <libpsio/psio.h>
#include <libiwl/iwl.h>
#include <libipv1/ip_lib.h>
#include <libqt/qt.h>
#include <math.h>
}

// define global variables for PSI3 libraries
extern "C" {
  char *psi_file_prefix;
  FILE *infile, *outfile;
  char *gprgid() { char *prgid = ":PSICHEM"; return(prgid); }
}

#endif

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Initializes the psi3 input/output system.
///
///////////////////////////////////////////////////////////////////////////////////////////////////

static void formic_init_psi3_io() {

  #ifndef FORMIC_HAVE_PSI3
  throw formic::Exception("cannot initialize psi3 io.  please recompile with psi3.");
  #endif

  #ifdef FORMIC_HAVE_PSI3

  class Psi3ioHandler {

    public:

      Psi3ioHandler() {
        const int argc = 0;
        char ** argv;
        //const int errcod = psi_start(argc, argv, 0);
        const int errcod = psi_start(&infile, &outfile, &psi_file_prefix, argc, argv, 0);
        if ( errcod != PSI_RETURN_SUCCESS )
          throw formic::Exception( (boost::format("psi_start failed with code %i in formic_init_psi3_io") % errcod).str() );
        //ffile(&infile, "psichem.inp", 2);
        //ffile(&outfile, "psichem.out", 1);
        ip_cwk_add(gprgid());
        tstart(outfile);
        psio_init();
        psio_ipv1_config();
        chkpt_init(PSIO_OPEN_OLD);
      }

      ~Psi3ioHandler() {
        chkpt_close();
        psio_done();
        //psi_stop();
        psi_stop(infile, outfile, psi_file_prefix);
      }

  };

  // The first time this function is called, the handler's constructor initializes the io system.
  // At program end, the handler's destructor finalizes the io system.
  static Psi3ioHandler handler;

  #endif

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   reads the number of orbitals in each irreducible representation
///
/// \param[in]   userinp  user input options
/// \param[out]  opirr    on exit, the number of orbitals in each irrep
///
///////////////////////////////////////////////////////////////////////////////////////////////////

void formic::get_psi3_orbs_per_irrep(const formic::InputBase & userinp, std::vector<int> & opirr) {

  #ifndef FORMIC_HAVE_PSI3
  throw formic::Exception("cannot read psi3 orbitals per irrep.  please recompile with psi3.");
  #endif

  #ifdef FORMIC_HAVE_PSI3

  // initialize psi3 io system
  formic_init_psi3_io();

  // get numbers of orbitals
  const int no = chkpt_rd_nso();

  // check that the number of orbitals matches the number of lattice sites
  if ( userinp.get<int>("nao") != no )
    throw formic::Exception("number of orbitals does not match psi3 orbitals:  no = %i, nao = %i") % no % userinp.get<int>("nao");

  // get the number of irreps
  const int nirrep = chkpt_rd_nirreps();

  // get a pointer to the number of orbitals per irrep
  const int * const opirr_ptr = chkpt_rd_sopi();

  // put the above information in the supplied vector
  opirr.resize(nirrep);
  std::copy(opirr_ptr, opirr_ptr + nirrep, opirr.begin());

  #endif

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   reads the one and two electron integrals from PSI3
///
/// \param[in]   userinp  user input options
/// \param[out]  enuc     the nuclear repulsion energy
/// \param[out]  ovi      the atomic orbital overlap integrals, size = no*no
/// \param[out]  oei      the one electron integrals, size = no*no
/// \param[out]  tei      the two electron integrals, size = no*(no+1) * ( no*(no+1) + 2 ) / 8
///
///////////////////////////////////////////////////////////////////////////////////////////////////

void formic::get_psi3_ints(const formic::InputBase & userinp, double & enuc, double * const ovi, double * const oei, double * const tei) {

  #ifndef FORMIC_HAVE_PSI3
  throw formic::Exception("cannot read psi3 integrals.  please recompile with psi3.");
  #endif

  #ifdef FORMIC_HAVE_PSI3

  // initialize psi3
  formic_init_psi3_io();

  // load the nuclear repulsion energy
  enuc = chkpt_rd_enuc();

  // get numbers of orbitals
  const int no = chkpt_rd_nso();       // total
  //const int nc = userinp.n_cor_orb();  // core
  //const int na = userinp.n_act_orb();  // active

  // check that the number of orbitals matches the number of lattice sites
  if ( userinp.get<int>("nao") != no )
    throw formic::Exception("number of orbitals does not match psi3 orbitals:  no = %i, nao = %i") % no % userinp.get<int>("nao");

//  // check that the numbers of orbitals are sane
//  if ( nc + na > no )
//    throw formic::Exception("too many core and/or active orbitals.  nc = %i, na = %i, no = %i") % nc % na % no;
//
//  // check that the number of active orbitals matches the number of lattice sites
//  if ( userinp.nsites() != 2*na )
//    throw formic::Exception("number of active orbitals must equal lattice sites:  na = %i, nsites = %i") % na % (userinp.nsites()/2);

//  // create temporary vectors into which we will load the integrals
//  std::vector<double> temp_int1(no*no, 0.0);
//  std::vector<double> temp_int2( no*(no+1)/2 * ( no*(no+1)/2 + 1 ) / 2, 0.0);
//
//  std::vector<double> ao_overlap(no*no, 0.0);

  // zero the integral arrays
  std::fill(ovi, ovi + no*no, 0.0);
  std::fill(oei, oei + no*no, 0.0);
  std::fill(tei, tei + ( no*(no+1) * ( no*(no+1) + 2 ) / 8 ), 0.0);

  // get the number of elements in the one electron integral's lower triangle
  const int ntri = no*(no+1)/2;

  // create a vector used in loading one electron data from PSI3
  std::vector<double> lower_triangle(ntri, 0.0);

  // load kinetic energy integrals
  lower_triangle.assign(lower_triangle.size(), 0.0);
  iwl_rdone(PSIF_OEI, PSIF_SO_T, &lower_triangle.at(0), ntri, 0, 0, outfile);
  formic::unpack_psi3_lower_triangle(no, &lower_triangle.at(0), oei);

  // load nuclear attraction integrals
  lower_triangle.assign(lower_triangle.size(), 0.0);
  iwl_rdone(PSIF_OEI, PSIF_SO_V, &lower_triangle.at(0), ntri, 0, 0, outfile);
  formic::unpack_psi3_lower_triangle(no, &lower_triangle.at(0), oei);

  // load orbital overlap integrals
  lower_triangle.assign(lower_triangle.size(), 0.0);
  iwl_rdone(PSIF_OEI, PSIF_SO_S, &lower_triangle.at(0), ntri, 0, 0, outfile);
  formic::unpack_psi3_lower_triangle(no, &lower_triangle.at(0), ovi);

//  // load the one electron integrals in lower triangle format
//  if ( userinp.get<bool>("s_half_orth") ) {
//    lower_triangle.assign(lower_triangle.size(), 0.0);
//    iwl_rdone(PSIF_OEI, PSIF_SO_T, &lower_triangle.at(0), ntri, 0, 0, outfile);
//    formic::unpack_psi3_lower_triangle(no, &lower_triangle.at(0), &temp_int1.at(0));
//    lower_triangle.assign(lower_triangle.size(), 0.0);
//    iwl_rdone(PSIF_OEI, PSIF_SO_V, &lower_triangle.at(0), ntri, 0, 0, outfile);
//    formic::unpack_psi3_lower_triangle(no, &lower_triangle.at(0), &temp_int1.at(0));
//    lower_triangle.assign(lower_triangle.size(), 0.0);
//    iwl_rdone(PSIF_OEI, PSIF_SO_S, &lower_triangle.at(0), ntri, 0, 0, outfile);
//    formic::unpack_psi3_lower_triangle(no, &lower_triangle.at(0), &ao_overlap.at(0));
//  } else {
//    iwl_rdone(PSIF_OEI, PSIF_MO_OEI, &lower_triangle.at(0), ntri, 0, 0, outfile);
//    formic::unpack_psi3_lower_triangle(no, &lower_triangle.at(0), &temp_int1.at(0));
//  }

  // precompute offset array used in loading two electron integrals
  std::vector<int> ioff(ntri*ntri, 0);
  for (int i = 0; i < ioff.size(); i++)
    ioff[i] = i*(i+1)/2;

  // load two electron integrals
  iwl_rdtwo(PSIF_SO_TEI, tei, &ioff.at(0), no, 0, 0, 0, outfile);
//  if ( userinp.get<bool>("s_half_orth") )
//    iwl_rdtwo(PSIF_SO_TEI, &temp_int2.at(0), &ioff.at(0), no, 0, 0, 0, outfile);
//  else
//    iwl_rdtwo(PSIF_MO_TEI, &temp_int2.at(0), &ioff.at(0), no, 0, 0, 0, outfile);

//  // perform S^(-1/2) orthogonalization if requested
//  if ( userinp.get<bool>("s_half_orth") )
//    formic::orthogonalize_orbs_by_overlap(no, ao_overlap, temp_int1, temp_int2);
//
//  // extract active space one electron integrals
//  for (int p = 0; p < na; p++)
//  for (int q = 0; q < na; q++) {
//    oei[na*p+q] = temp_int1[no*(p+nc)+(q+nc)];
//    for (int i = 0; i < nc; i++)
//      oei[na*p+q] +=   2.0 * temp_int2[formic::Hamiltonian<double>::tei_compound_index(i, i, p+nc, q+nc)]
//                     - 1.0 * temp_int2[formic::Hamiltonian<double>::tei_compound_index(i, p+nc, i, q+nc)];
//  }
//
//  // extract active space two electron integrals
//  for (int p = 0; p < na; p++)
//  for (int q = 0; q <= p; q++)
//  for (int r = 0; r < na; r++)
//  for (int s = 0; s <= r; s++) {
//    const int tmp_cmpd = formic::Hamiltonian<double>::tei_compound_index(p+nc, q+nc, r+nc, s+nc);
//    const int act_cmpd = formic::Hamiltonian<double>::tei_compound_index(p, q, r, s);
//    tei[act_cmpd] = temp_int2[tmp_cmpd];
//  }
//
//  // add core orbital contribution to the nuclear energy (nuclear energy is used to also hold core energy)
//  for (int i = 0; i < nc; i++)
//    enuc += 2.0 * temp_int1[no*i+i] + temp_int2[formic::Hamiltonian<double>::tei_compound_index(i, i, i, i)];
//  for (int i = 0; i < nc; i++)
//  for (int j = 0; j <  i; j++)
//    enuc +=   4.0 * temp_int2[formic::Hamiltonian<double>::tei_compound_index(i, i, j, j)]
//            - 2.0 * temp_int2[formic::Hamiltonian<double>::tei_compound_index(i, j, i, j)];

  #endif

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   throws an error because PSI3 doesn't use complex integrals
///
/// \param[in]   userinp  user input options
/// \param[out]  enuc     the nuclear repulsion energy
/// \param[out]  ovi      the atomic orbital overlap integrals, size = no*no
/// \param[out]  oei      the one electron integrals, size = na*na
/// \param[out]  tei      the two electron integrals, size = na*(na+1) * ( na*(na+1) + 2 ) / 8
///
///////////////////////////////////////////////////////////////////////////////////////////////////


void formic::get_psi3_ints(const formic::InputBase & userinp, double & enuc, std::complex<double> * const ovi,
                                                                             std::complex<double> * const oei,
                                                                             std::complex<double> * const tei) {

  throw formic::Exception("Cannot read complex integrals from PSI3.  Please use scalar_type = real");

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   expands the psi3 lower triangle storage for one electron integrals
///
///   This function adds the one electron integral data stored in the lower triangle to the one
///   electron integral data stored in the full matrix.
///
/// \param[in]      n        the number of spatial orbitals
/// \param[in]      lt       the one electron integrals in lower triangle storage, size n*(n+1)/2
/// \param[in,out]  oei      the one electron integrals in full-matrix storage, size = n*n
///
///////////////////////////////////////////////////////////////////////////////////////////////////

void formic::unpack_psi3_lower_triangle(const int n, const double * const lt, double * const oei) {

  for (int i = 0; i < n; i++) {

    // diagonal elements
    oei[i*n+i] += lt[ i*(i+1)/2 + i ];

    // off diagonal elements
    for (int j = 0; j < i; j++) {
      oei[i*n+j] += lt[ i*(i+1)/2 + j ];
      oei[j*n+i] += lt[ i*(i+1)/2 + j ];
    }

  }

}
